from AStock.ASQuery import (
    ASQuery_stock_info,
    ASQuery_research_report,
    ASQuery_stock_day,
    ASQuery_trade_cal
)
from AStock.ASSetting import settings
import argparse
import pandas as pd
from pymongo import UpdateOne
from datetime import datetime, timedelta


def _main_(args):
    begin_date = args.begin_date  # 研报开始日期
    end_date = args.end_date  # 研报终止日期
    close_date = args.close_date  # 收盘价日期

    # 如果收盘价日期休市或者还未收盘，找前面最近的开市一天
    now = datetime.now()
    now_date = now.strftime('%Y%m%d')
    if now_date == close_date:
        if now.hour < 15:
            cdate = datetime.strptime(close_date, '%Y%m%d')
            cdate -= timedelta(days=1)
            close_date = cdate.strftime('%Y%m%d')
    while not ASQuery_trade_cal(close_date):
        cdate = datetime.strptime(close_date, '%Y%m%d')
        cdate -= timedelta(days=1)
        close_date = cdate.strftime('%Y%m%d')
    print('real close date: {}'.format(close_date))

    # 获取股票信息
    df_stock_info = ASQuery_stock_info(fields=['code'])

    codes = df_stock_info['code'].tolist()

    # 获取当前收盘价
    df_close = ASQuery_stock_day(codes, close_date, close_date, fields=['code', 'close'])

    # 获取机构研报盈利预测（最近三个月最近最多10只研报预测盈利均值）
    predict_list = []
    for i, code in enumerate(codes):
        print('get research report {} {}/{}'.format(code, i, len(codes)))
        df_research_report = ASQuery_research_report(
            code,
            begin_date,
            end_date,
            fields=['publishDate', 'predictThisYearEps', 'predictNextYearEps', 'predictNextTwoYearEps', 'count'],
            limit=10
        )
        # print(df_research_report.head(10))
        if not df_research_report.empty:
            predict_dict = {
                'code': code,
                'predictThisYearEps': df_research_report['predictThisYearEps'].dropna().mean(),
                'predictNextYearEps': df_research_report['predictNextYearEps'].dropna().mean(),
                'predictNextTwoYearEps': df_research_report['predictNextTwoYearEps'].dropna().mean(),
                'recentMonthReportCount': df_research_report.loc[0, 'count']
            }
        else:
            predict_dict = {
                'code': code,
                'predictThisYearEps': None,
                'predictNextYearEps': None,
                'predictNextTwoYearEps': None,
                'recentMonthReportCount': 0
            }
        predict_list.append(predict_dict)

    df = pd.DataFrame(predict_list)

    df = df.join(df_close.set_index('code'), on='code', how='left')
    df['close_date'] = int(close_date)

    # 计算PE和PEG
    df = df.assign(
        predictThisYearPe=df.close / df.predictThisYearEps,
        predictNextYearPe=df.close / df.predictNextYearEps,
    )
    df = df.assign(
        predictThisYearPeg=df.predictThisYearPe / (100 * (df.predictNextYearEps - df.predictThisYearEps) / df.predictThisYearEps),
        predictNextYearPeg=df.predictNextYearPe / (100 * (df.predictNextTwoYearEps - df.predictNextYearEps) / df.predictNextYearEps)
    )
    df = df.round(2)

    # 存库
    # code
    # close
    # close_date
    # predictThisYearEps
    # predictNextYearEps
    # predictThisYearPe
    # predictNextYearPe
    # recentMonthReportCount
    bulk = []
    for row in df.itertuples():
        values = row._asdict()
        del values['Index']
        update_one = UpdateOne(
            {'code': row.code},
            {'$set': values},
            upsert=True
        )
        bulk.append(update_one)

    try:
        coll = settings.database.stock_valuation
        coll.bulk_write(bulk)
    except Exception as e:
        print(e)
    print('saved {} items to stock_valuation'.format(len(bulk)))


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='save valuation')
    argparser.add_argument('-b', '--begin-date', required=True,
                           help='begin date of research report, yyyymmdd format')
    argparser.add_argument('-e', '--end-date', required=True,
                           help='end date of research report, yyyymmdd format')
    argparser.add_argument('-c', '--close-date', required=True,
                           help='close date of price, yyyymmdd format')
    args = argparser.parse_args()
    _main_(args)

